#include "BVH.h"

#include <queue>
using namespace std;

#include <boost/assert.hpp>

BV_AABB::BV_AABB( const Vec3f& centerOrMax, const Vec3f& radiusOrMin, const ConstructionMode mode /*= CenterRadiusConstruction*/ )
{
    Update(centerOrMax,radiusOrMin, mode);
}

void BV_AABB::Update( const Vec3f& centerOrMax, const Vec3f& radiusOrMin, const ConstructionMode mode /*= CenterRadiusConstruction*/ )
{
    switch (mode) {
            case CenterRadiusConstruction:
                _center = centerOrMax;
                _radius = radiusOrMin;
                break;
            case MaxMinConstruction:
                _center = (centerOrMax + radiusOrMin) / 2;
                _radius = (centerOrMax - radiusOrMin);
                break;
            default: break;
    }
}

llib::Boolean BVNode::BVIntersection( const BVNode* another ) const
{
    return CheckBVIntersection(this, another);
}

llib::Boolean BVNode::CheckBVIntersection( const BVNode* left, const BVNode* right )
{
    Vec3f lmax = left->_bv.GetMax();
    Vec3f lmin = left->_bv.GetMin();
    Vec3f rmax = right->_bv.GetMax();
    Vec3f rmin = right->_bv.GetMin();

    return !( lmin.x > rmax.x || lmax.x < rmin.x ||
              lmin.y > rmax.y || lmax.y < rmin.y ||
              lmin.z > rmax.z || lmax.z < rmin.z );
}

void BVNode::UpdateBV( const Vec3f& centerOrMax, const Vec3f& radiusOrMin, const BV_AABB::ConstructionMode mode /*= BV_AABB::CenterRadiusConstruction*/ )
{
    _bv.Update(centerOrMax, radiusOrMin, mode);
}

BVTree::~BVTree()
{
    delete[] _nodeList;
    delete[] _leafElemIndexList;
}

BVTree::BVTree( const BVTree& tree )
{
    copy(tree);
}

BVTree& BVTree::operator=( const BVTree& tree )
{
    // pointer are NULLs when empty
    delete[] _nodeList;
    delete[] _leafElemIndexList;
    copy(tree);
    return *this;
}

// We build a complete binary tree as BVH
// thus tree can be arranged in array
void BVTree::Build( const vector<Vec3f>& elementPostions )
{
    BOOST_ASSERT(SZ(elementPostions) > 0);

    // remove old data
    delete[] _nodeList;
    delete[] _leafElemIndexList;

    /* index pointer to elementPostions */
    vector<Int32> leafIndex(SZ(elementPostions));
    FOR(I,0,SZ(leafIndex)) leafIndex[I] = I;
    /* arrange elements by index pointer */
    leafArrangement(elementPostions, leafIndex);

    /* the node numbers */
    _leafNum = SZ(elementPostions);
    _nodeNum = _leafNum * 2 - 1;
    /* allocate */
    _nodeList = new BVNode[_nodeNum];
    _leafElemIndexList = new Int32[_leafNum];
    /* first and last node index */
    Int32 power = 0, tmp = _leafNum - 1;
    while(tmp > 0) {
        tmp >>= 1;
        ++power;
    }
    Int32 maxLeafNum = 1 << power;
    _firstLeafIndex = (1 << power) - 1;
    Int32 minLeafIndex = _nodeNum >> 1; // the minimum index leaf is in middle of array
    if(minLeafIndex == _firstLeafIndex) ++power;
    _lastLeafIndex = (1 << power) - 2;
    /* place the leaf element into _leafElemIndexList */
    Int32 curIndex = _leafNum * 2 - maxLeafNum;
    if(_leafNum == maxLeafNum) curIndex = 0;
    FOR(I,0,_leafNum) {
        _leafElemIndexList[I] = leafIndex[curIndex++];
        if(curIndex >= _leafNum)
            curIndex = 0;
    }

}

BVNode* BVTree::GetLeftChild( const BVNode* node ) const
{
    Int32 index = getNodeIndex(node);
    BOOST_ASSERT(index != -1);
    Int32 leftChildIndex = index * 2 + 1;
    if(leftChildIndex > _nodeNum - 1) return NULL;
    return _nodeList + leftChildIndex;
}

BVNode* BVTree::GetRightChild( const BVNode* node ) const
{
    Int32 index = getNodeIndex(node);
    BOOST_ASSERT(index != -1);
    if(index == -1) return NULL;
    Int32 rightChildIndex = index * 2 + 2;
    if(rightChildIndex > _nodeNum - 1) return NULL;
    return _nodeList + rightChildIndex;
}

llib::Boolean BVTree::IsLeafNode( const BVNode* node ) const
{
    Int32 index = getNodeIndex(node);
    BOOST_ASSERT(index != -1);
    return index * 2 + 1 > _nodeNum - 1;
}

llib::Int32 BVTree::GetLeafElemIndex( const BVNode* node ) const
{
    BOOST_ASSERT(IsLeafNode(node));
    Int32 leafIndex = getLeafIndex(node);
    BOOST_ASSERT(leafIndex != -1);
    return _leafElemIndexList[leafIndex];
}

void BVTree::TraversalBegin( TraversalType type )
{
    switch (type) {
        case Preorder:  traversalBegin_Preorder();  break;
        case Inorder:   traversalBegin_Inorder();   break;
        case Postorder: traversalBegin_Postorder(); break;
        case Hierarchy: BOOST_ASSERT(FALSE);        break;  // not impl yet
        case LeafOnly:  BOOST_ASSERT(FALSE);        break;  // not impl yet
        default:        BOOST_ASSERT(FALSE);        break;  // never get here
    }
}

BVNode* BVTree::TraversalNext()
{
    return (this->*_traversalNextFunPtr)();
}

void BVTree::Traversal( TraversalType type, Boolean (*callback)(const BVTree& tree, BVNode* curNode) )
{
    TraversalBegin(type);
    BVNode* node = TraversalNext();
    while(node != NULL) {
        if( !(*callback)(*this, node) ) break;
        node = TraversalNext();
    }
    TraversalEnd();
}

template <class Functor_t>
void BVTree::Traversal( TraversalType type, Functor_t& functor )
{
    TraversalBegin(type);
    BVNode* node = TraversalNext();
    while(node != NULL) {
        if( !functor(*this, node) ) break;
        node = TraversalNext();
    }
    TraversalEnd();
}

void BVTree::copy( const BVTree& tree )
{
    if(tree._nodeNum != 0 && tree._leafNum != 0) {
        _nodeNum            = tree._nodeNum;
        _leafNum            = tree._leafNum;
        _firstLeafIndex     = tree._firstLeafIndex;
        _lastLeafIndex      = tree._lastLeafIndex;
        _nodeList           = new BVNode[_nodeNum];
        _leafElemIndexList  = new Int32[_leafNum];
        FOR(I,0,_nodeNum)
            _nodeList[I] = tree._nodeList[I];
        FOR(I,0,_leafNum)
            _leafElemIndexList[I] = tree._leafElemIndexList[I];
    } else {
        _nodeNum            = 0;
        _leafNum            = 0;
        _firstLeafIndex     = -1;
        _lastLeafIndex      = -1;
        _nodeList           = NULL;
        _leafElemIndexList  = NULL;
    }
}

llib::Int32 BVTree::getPrevLeafIndex( const Int32 nodeId ) const
{
    if(nodeId == _firstLeafIndex) return -1;
    Int32 minLeafIndex = _leafNum / 2;
    if(nodeId == minLeafIndex) return _nodeNum - 1;
    return nodeId - 1;
}

llib::Int32 BVTree::getNextLeafIndex( const Int32 nodeId ) const
{
    if(nodeId == _lastLeafIndex) return -1;
    Int32 minLeafIndex = _leafNum / 2;
    if(nodeId == _nodeNum - 1) return minLeafIndex;
    return nodeId + 1;
}

/* Arrange element by iterative division on longest axis,
 * result are saved in the order index list leafIndex,
 * in the form of leaf nodes of the corresponding complete binary tree.
 */
void BVTree::leafArrangement( const vector<Vec3f>& elementPostions, vector<Int32>& leafIndex ) const
{
#define ELEMENT_POSITION(id)    elementPostions[ leafIndex[id] ]

    typedef pair<Int32,Int32> Range;
    queue<Range> callHeap;  // here we're using the nonrecursive mothed
    Range r(0,SZ(elementPostions)), r2;
    callHeap.push(r);
    Int32 rangeLen = 0;
    Int8  maxAxis = 0;
    Vec3f boundMax, boundMin;
    const static Int8 axisSelect[8] = { 2, 2, -1, 0, 1, -1, 1, 0 };
    Int32 tmp, division, pivotId, pivotVal;
    Float divideVal;
    Int32 lid, rid;

    while(!callHeap.empty()) {
        r = callHeap.front();
        callHeap.pop();
        rangeLen = r.second - r.first;
        if(rangeLen < 3) {
            if(rangeLen == 2) {
                if(ELEMENT_POSITION(r.first) > ELEMENT_POSITION(r.first+1))
                    swap(leafIndex[r.first],leafIndex[r.first+1]);
            }
            continue; // no need to go deep
        }

        // max bound axis selection
        boundMax = boundMin = ELEMENT_POSITION(r.first);
        FOR(pid,r.first+1,r.second) {
            const Vec3f& pos = ELEMENT_POSITION(pid);
            FOR(axis,0,3) {
                boundMin[axis] = min(boundMin[axis], pos[axis]);
                boundMax[axis] = max(boundMax[axis], pos[axis]);
            }
        }
        boundMax -= boundMin; // boundMax now saves boundRange
        tmp = 0;
        if(boundMax.x > boundMax.y) tmp |= 1;
        if(boundMax.x > boundMax.z) tmp |= 2;
        if(boundMax.y > boundMax.z) tmp |= 4;
        maxAxis = axisSelect[tmp];

        // arrangement
        // we're using the O(n) algorithm to reorder array into two part,
        // the smaller and bigger parts.
        division = completeBinaryTreeLeafDivide(rangeLen);
        division += r.first;
        r2 = r;
        while(TRUE) {
            pivotId = math::RndRange(r2.first, r2.second);
            //pivotId = r2.first; 
            pivotVal = leafIndex[pivotId];
            divideVal = ELEMENT_POSITION(pivotId)[maxAxis];
            leafIndex[pivotId] = leafIndex[r2.first];
            lid = r2.first;
            rid = r2.second - 1;
            while(lid < rid) {
                while(lid < rid && ELEMENT_POSITION(rid)[maxAxis] >= divideVal) --rid;
                leafIndex[lid] = leafIndex[rid];
                while(lid < rid && ELEMENT_POSITION(lid)[maxAxis] <= divideVal) ++lid;
                leafIndex[rid] = leafIndex[lid];
            }
            leafIndex[lid] = pivotVal;
            if(lid == division) break;
            else if(lid > division) r2.second = lid;
            else r2.first = lid + 1;
        };

        // the sub tree process
        r2.first = r.first;
        r2.second = division;
        callHeap.push(r2);
        r2.first = division;
        r2.second = r.second;
        callHeap.push(r2);
    }

#undef ELEMENT_POSITION
}

/* Divide leaf into two part (the left and right sub tree),
 * according to the arrangement of complete binary tree.
 * leafNum : the tree leaf number
 * return value: the left sub tree leaf number
 */
llib::Int32 BVTree::completeBinaryTreeLeafDivide( Int32 leafNum ) const
{
    // the full binary tree case
    if((leafNum & leafNum - 1) == 0) return leafNum / 2;
    // the other (common) case
    int childMax = 1;
    while((childMax << 1) < leafNum) childMax <<= 1;
    int childMin = childMax >> 1;
    int res = leafNum - childMin;
    if(res > childMax) res = childMax;
    return res;
}

BVNode* BVTree::traversalNext_Preorder()
{
    if(_traversalHelper.empty()) return NULL;

    Int32 nodeId = _traversalHelper.back();
    _traversalHelper.pop_back();
    BVNode* res = GetRoot() + nodeId;

    Int32 childId = nodeId * 2 + 2;
    if(childId < _nodeNum) _traversalHelper.push_back(childId);
    --childId;
    if(childId < _nodeNum) _traversalHelper.push_back(childId);

    return res;
}

BVNode* BVTree::traversalNext_Inorder()
{
    if(_traversalHelper.empty()) return NULL;

    Int32 nodeId = _traversalHelper.back();
    _traversalHelper.pop_back();
    BVNode* res = GetRoot() + nodeId;

    nodeId = nodeId * 2 + 2;
    while(nodeId < _nodeNum) {
        _traversalHelper.push_back(nodeId);
        nodeId = nodeId * 2 + 1;
    }

    return res;
}

BVNode* BVTree::traversalNext_Postorder()
{
    if(_traversalHelper.empty()) return NULL;

    BVNode* res = NULL;
    Int32 nodeId, rChildId;
    while(res == NULL) {
        nodeId = _traversalHelper.back();
        rChildId = nodeId * 2 + 2;
        if(rChildId >= _nodeNum || rChildId == _lastPop) {
            _traversalHelper.pop_back();
            res = GetRoot() + nodeId;
            _lastPop = nodeId;
        } else {
            nodeId = rChildId;
            while(nodeId < _nodeNum) {
                _traversalHelper.push_back(nodeId);
                nodeId = nodeId * 2 + 1;
            }
        }
    }

    return res;
}

